#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>
#include "common.h"

#define LOCAL_ROUTES_MAX_LENGTH 1024
#define SRV6_SID_CONFIGS_MAX_LENGTH 1024

char _license[] SEC("license") = "GPL";

struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, DeviceType);
} config_device_type SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __uint(max_entries, LOCAL_ROUTES_MAX_LENGTH);
    __uint(map_flags, BPF_F_NO_PREALLOC);
    __type(key, Ipv4LpmKey);
    __type(value, Srv6Route);
} config_ipv4_routes SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __uint(max_entries, LOCAL_ROUTES_MAX_LENGTH);
    __uint(map_flags, BPF_F_NO_PREALLOC);
    __type(key, Ipv6LpmKey);
    __type(value, Srv6Route);
} config_ipv6_routes SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, SRV6_SID_CONFIGS_MAX_LENGTH);
    __type(key, Ipv6Addr8);
    __type(value, LocalSrv6SidConfig);
} config_local_sid_configs SEC(".maps");


static int fib_lookup_src_ipv6(struct __sk_buff* skb, void* src_ipv6, void* dst_ipv6) {
    struct bpf_fib_lookup fib_params = {0};
    fib_params.family = AF_INET6;
    fib_params.ifindex = skb->ingress_ifindex;
    __builtin_memcpy(fib_params.ipv6_dst, dst_ipv6, IPV6_ADDR_SIZE);
    int ret = bpf_fib_lookup(skb, &fib_params, sizeof(fib_params), BPF_FIB_LOOKUP_SKIP_NEIGH | BPF_FIB_LOOKUP_OUTPUT | BPF_FIB_LOOKUP_SRC);
    if (ret != 0) {
        return ret;
    }
    __builtin_memcpy(src_ipv6, fib_params.ipv6_src, IPV6_ADDR_SIZE);
    return 0;
}

static int srv6_encap(struct __sk_buff* skb, __u32 ip_offset, const Srv6Route* route, __u8 srv6_srh_next_header, __be16 original_total_len) {
    __u32 srv6_segment_list_length = 0;
    for (__u32 i = 0; i < SRV6_SEGMENT_LIST_MAX_LENGTH; i++) {
        if (ipv6_addr8_is_zero(&route->segments[i])) {
            break;
        }
        srv6_segment_list_length++;
    }

    if (srv6_segment_list_length == 0 || srv6_segment_list_length > SRV6_SEGMENT_LIST_MAX_LENGTH) {
        return TC_ACT_SHOT;
    }

    __u32 srh_len = sizeof(SegmentRoutingHeader) + srv6_segment_list_length * IPV6_ADDR_SIZE;

    // 增加头部空间
    if (bpf_skb_adjust_room(skb, IPV6_HDR_LEN + srh_len, BPF_ADJ_ROOM_MAC, BPF_F_ADJ_ROOM_ENCAP_L3_IPV6) != 0) {
        return TC_ACT_SHOT;
    }

    void* data_end = (void*)(long)skb->data_end;

    struct ipv6hdr* new_iph = (struct ipv6hdr*)(((void*)(long)skb->data) + ip_offset);
    if ((void*)(new_iph + 1) > data_end) {
        return TC_ACT_SHOT;
    }

    __builtin_memset(new_iph, 0, sizeof(struct ipv6hdr));

    new_iph->version = 6;
    new_iph->priority = 0;
    new_iph->flow_lbl[0] = 0;
    new_iph->flow_lbl[1] = 0;
    new_iph->flow_lbl[2] = 0;
    new_iph->payload_len = bpf_htons(bpf_ntohs(original_total_len) + srh_len);
    new_iph->nexthdr = IPPROTO_ROUTING;
    new_iph->hop_limit = 64;

    // 设置源地址
    if (fib_lookup_src_ipv6(skb, new_iph->saddr.in6_u.u6_addr8, route->segments[srv6_segment_list_length - 1].addr8) !=0) {
        return TC_ACT_SHOT;
    }

    // 设置目标地址为第一个SRv6段
    __builtin_memcpy(new_iph->daddr.in6_u.u6_addr8, route->segments[srv6_segment_list_length - 1].addr8, IPV6_ADDR_SIZE);

    // 构建SRH
    SegmentRoutingHeader *srh = (SegmentRoutingHeader *)(new_iph + 1);
    if ((void*)(srh + 1) > data_end) {
        return TC_ACT_SHOT;
    }

    srh->next_header = srv6_srh_next_header;
    srh->hdr_ext_len = (srh_len / 8) - 1;
    srh->routing_type = 4; // SRH类型
    srh->segments_left = srv6_segment_list_length - 1;
    srh->last_entry = srv6_segment_list_length - 1;
    srh->flags = 0;
    srh->tag = 0;

    Ipv6Addr8* segment_list = (Ipv6Addr8*)(srh + 1);
    if ((void*)(segment_list + srv6_segment_list_length + 1) > data_end) {
        return TC_ACT_SHOT;
    }

    __u32 i = 0;
    bpf_for(i, 0, srv6_segment_list_length) {
        if ((void*)(segment_list + i) + IPV6_ADDR_SIZE > data_end) {
            return TC_ACT_SHOT;
        }
        __builtin_memcpy(segment_list[i].addr8, route->segments[i].addr8, IPV6_ADDR_SIZE);
    }
    return TC_ACT_OK;
}

static int srv6_end(struct __sk_buff *skb, struct ipv6hdr *iph, SegmentRoutingHeader *srh) {
    void* data_end = (void*)(long)skb->data_end;
    const Ipv6Addr8* segment_list = (Ipv6Addr8*)(srh + 1);

    if (srh->segments_left == 0) {
        return TC_ACT_OK;
    }

    // 减少Segments Left
    srh->segments_left--;
    
    // 更新目标地址为下一个SID
    if (srh->segments_left < srh->last_entry + 1) {
        __u32 next_sid_index = srh->segments_left;

        // 确保访问的索引不会导致越界
        if (next_sid_index >= SRV6_SEGMENT_LIST_MAX_LENGTH) {
            return TC_ACT_SHOT;
        }
        
        // 确保要访问的地址在数据包范围内
        if ((void*)(segment_list + next_sid_index + 1) > data_end) {
            return TC_ACT_SHOT;
        }

        if (fib_lookup_src_ipv6(skb, iph->saddr.in6_u.u6_addr8, iph->daddr.in6_u.u6_addr8) != 0) {
            return TC_ACT_SHOT;
        }

        // 复制下一个SID到IPv6目标地址
        __builtin_memcpy(iph->daddr.in6_u.u6_addr8, 
                        segment_list[next_sid_index].addr8, 
                        IPV6_ADDR_SIZE);
    }
    
    return TC_ACT_OK;
}

static int srv6_end_dt46(struct __sk_buff *skb, SegmentRoutingHeader *srh) {
    void* data_end = (void*)(long)skb->data_end;

    // 计算SRH长度
    __u32 srh_len = (srh->hdr_ext_len + 1) * 8;
    
    // 检查SRH是否完整
    if ((void*)((char*)srh + srh_len) > data_end) {
        return TC_ACT_SHOT;
    }

    __u64 flags = 0;
    switch (srh->next_header) {
    case IPPROTO_IPIP:
        flags = BPF_F_ADJ_ROOM_DECAP_L3_IPV4;
        break;
    case IPPROTO_IPV6:
        flags = BPF_F_ADJ_ROOM_DECAP_L3_IPV6;
        break;
    default:
        return TC_ACT_SHOT;
    }
    
    // 移除外部IPv6头部和SRH
    if (bpf_skb_adjust_room(skb, -(IPV6_HDR_LEN + srh_len), BPF_ADJ_ROOM_MAC, flags) != 0) {
        return TC_ACT_SHOT;
    }

    return TC_ACT_OK;
}

static int srv6_handle(struct __sk_buff *skb, struct ipv6hdr *iph) {
    void* data_end = (void*)(long)skb->data_end;

    Ipv6Addr8 sid_key;
    __builtin_memcpy(&sid_key, iph->daddr.in6_u.u6_addr8, IPV6_ADDR_SIZE);

    // 查找SID配置
    const LocalSrv6SidConfig *local_sid_config = bpf_map_lookup_elem(&config_local_sid_configs, &sid_key);
    if (!local_sid_config) {
        return TC_ACT_OK;
    }
    
    // 检查是否是SRv6数据包
    if (iph->nexthdr != IPPROTO_ROUTING) {
        return TC_ACT_SHOT;
    }

    SegmentRoutingHeader *srh = (SegmentRoutingHeader *)(iph + 1);
    if ((void*)(srh + 1) > data_end) {
        return TC_ACT_SHOT;
    }
    
    // 检查是否是SRH (Routing Type = 4)
    if (srh->routing_type != 4) {
        return TC_ACT_SHOT;
    }

    switch (local_sid_config->behavior) {
    case SRV6_BEHAVIOR_END:
        return srv6_end(skb, iph, srh);
    case SRV6_BEHAVIOR_END_DT46:
        return srv6_end_dt46(skb, srh);
    default:
        return TC_ACT_OK;
    }
}

static int ipv4_handle(struct __sk_buff *skb, __u32 ip_offset, struct iphdr *iph) {
    Ipv4LpmKey key = {
        .prefix_length = 32,
    };

    __builtin_memcpy(key.address.addr8, &iph->daddr, IPV4_ADDR_SIZE);

    const Srv6Route* route = bpf_map_lookup_elem(&config_ipv4_routes, &key);
    if (!route) {
        return TC_ACT_OK;
    }

    return srv6_encap(skb, ip_offset, route, IPPROTO_IPIP, iph->tot_len);
}

static int ipv6_handle(struct __sk_buff *skb, __u32 ip_offset, struct ipv6hdr *iph) { 
    Ipv6LpmKey key = {
        .prefix_length = 128
    };

    __builtin_memcpy(key.address.addr8, iph->daddr.in6_u.u6_addr8, IPV6_ADDR_SIZE);

    const Srv6Route* route = bpf_map_lookup_elem(&config_ipv6_routes, &key);
    if (!route) {
        if (iph->nexthdr == IPPROTO_ROUTING) {
            return srv6_handle(skb, iph);
        } else {
            return TC_ACT_OK;
        }
    }
    return srv6_encap(skb, ip_offset, route, IPPROTO_IPV6, iph->payload_len + IPV6_HDR_LEN);
}


SEC("tc")
int wanmesh_srv6_router(struct __sk_buff *skb) {
    void *data_end = (void*)(long)skb->data_end;

    __u32 key = 0;
    DeviceType* device_type  = bpf_map_lookup_elem(&config_device_type, &key);
    if (!device_type) {
        return TC_ACT_OK;
    }

    __u32 ip_offset = 0;
    switch (*device_type){
    case DEVICE_TYPE_ETHERNET:
        ip_offset = IP_OFFSET_ETHERNET;
        break;
    case DEVICE_TYPE_RAW_IP:
        ip_offset = IP_OFFSET_RAW_IP;
        break;
    default:
        return TC_ACT_OK;
    }

    switch (skb->protocol) {
    case bpf_htons(ETH_P_IP): {
        struct iphdr* iph = (struct iphdr*)(((void*)(long)skb->data) + ip_offset);
        if ((void*)(iph + 1) > data_end) {
            return TC_ACT_SHOT;
        }
        return ipv4_handle(skb, ip_offset, iph);
    }
    case bpf_htons(ETH_P_IPV6): {
        struct ipv6hdr* iph = (struct ipv6hdr*)(((void*)(long)skb->data) + ip_offset);
        if ((void*)(iph + 1) > data_end) {
            return TC_ACT_SHOT;
        }
        return ipv6_handle(skb, ip_offset, iph);
    }
    default:
        return TC_ACT_OK;
    }

}